import pymysql
from pymysql.cursors import DictCursor
import os
from dotenv import load_dotenv

class DBUtils:
    def __init__(self, use_env=True):
        # 加载环境变量
        load_dotenv()
        
        # 记录是否使用环境变量配置
        self.use_env = use_env

        if use_env:
            # 从环境变量获取数据库配置
            self.host = os.getenv('DB_HOST', 'localhost')
            self.user = os.getenv('DB_USER', 'root')
            self.password = os.getenv('DB_PASSWORD', '')
            self.database = os.getenv('DB_NAME', 'ybp_zhxs')
            self.port = int(os.getenv('DB_PORT', '3306'))
            
            # 初始化数据库连接
            self.connect()
        else:
            # 使用book_crawler数据库的默认配置
            self.host = '127.0.0.1'
            self.user = 'root'
            self.password = '123456'
            self.database = 'book_crawler'
            self.port = 3306
            
            # 初始化数据库连接
            self.open_connect()

    def connect(self):
        """建立数据库连接"""
        try:
            self.conn = pymysql.connect(
                host=self.host,
                user=self.user,
                password=self.password,
                database=self.database,
                port=self.port,
                charset='utf8mb4',
                cursorclass=DictCursor
            )
            self.cursor = self.conn.cursor()
        except Exception as e:
            print(f"数据库连接失败: {e}")
            raise e
            
    def open_connect(self):
        """建立book_crawler数据库连接"""
        try:
            self.conn = pymysql.connect(
                host='127.0.0.1',
                port=3306,
                user='root',
                password='123456',
                database='book_crawler',
                charset='utf8mb4',
                cursorclass=DictCursor
            )
            self.cursor = self.conn.cursor()
            print("book_crawler数据库连接成功")
        except Exception as e:
            print(f"book_crawler数据库连接失败: {e}")
            raise e

    def execute_query(self, sql, params=None, fetch_one=False):
        """执行SQL查询"""
        try:
            # 如果连接已关闭，重新连接
            if not self.conn or not self.conn.open:
                if self.use_env:
                    self.connect()
                else:
                    self.open_connect()
            
            # 执行SQL查询
            self.cursor.execute(sql, params)

            # 获取结果
            if fetch_one:
                return self.cursor.fetchone()
            return self.cursor.fetchall()

        except Exception as e:
            print(f"查询执行失败: {e}")
            # 如果是连接错误，尝试重新连接
            if "MySQL Connection not available" in str(e):
                if self.use_env:
                    self.connect()
                else:
                    self.open_connect()
                return self.execute_query(sql, params, fetch_one)
            raise e

    def execute_update(self, sql, params=None):
        """执行更新操作（INSERT/UPDATE/DELETE）"""
        try:
            # 如果连接已关闭，重新连接
            if not self.conn or not self.conn.open:
                if self.use_env:
                    self.connect()
                else:
                    self.open_connect()

            # 执行SQL更新
            self.cursor.execute(sql, params)
            self.conn.commit()

            return self.cursor.rowcount

        except Exception as e:
            print(f"更新执行失败: {e}")
            if self.conn and self.conn.open:
                self.conn.rollback()
            # 如果是连接错误，尝试重新连接
            if "MySQL Connection not available" in str(e):
                if self.use_env:
                    self.connect()
                else:
                    self.open_connect()
                return self.execute_update(sql, params)
            raise e

    def close_connect(self):
        """关闭数据库连接"""
        try:
            if self.cursor:
                self.cursor.close()
            if self.conn and self.conn.open:
                self.conn.close()
        except Exception as e:
            print(f"关闭连接失败: {e}")
